from core.decorators.database_decorator import with_async_session
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, asc, delete
from core.utils.paginator import Paginator
from core.dependencies import AsyncSessionLocal

from src.modules.knowbase.repository.entity.vectors import ImageVector
from src.modules.knowbase.schemas.model_schema import ImageVectorResponse
from typing import List

class ImageVectorDao:
    @staticmethod
    @with_async_session()
    async def get_imagevector_list(
          user_id: str = None,  file_id:str = None,
          kb_id: str = None,
          page:int = 1,page_size:int=10,
          session :AsyncSession = None
    ):
        query = select(ImageVector)
        if user_id:
            query = query.where(ImageVector.user_id == user_id)
        if file_id:
            query = query.where(ImageVector.file_id == file_id)
        if kb_id:
            query = query.where(ImageVector.know_base_id == file_id)
        query = query.order_by(asc(ImageVector.uid))
        paginator = Paginator(session = session,
                              model = ImageVector,
                              response_model= ImageVectorResponse,
                              base_query= query)
        return await paginator.paginate(page, page_size)


    @staticmethod
    async def get_imagevector_list_for_rag(question_vector: List[float],
                                           session: AsyncSession,
                                           knowbase_ids: List[str],
                                           top_k: int):
        if not knowbase_ids:
            return []

        query = (
            select(
                ImageVector,
                ImageVector.desc_vector.l2_distance(question_vector).label('distance')
            )
            .where(ImageVector.know_base_id.in_(knowbase_ids))
            .order_by('distance')
            .limit(top_k)
        )
        result = await session.execute(query)

        image_results = result.all()
        print("all")
        print(image_results)
        return image_results

    @staticmethod
    async def delete_vector(vector_id:str):
        session:AsyncSession = AsyncSessionLocal()
        try:
            delete_sql = delete(ImageVector).where(ImageVector.vector_id == vector_id)
            result = await session.execute(delete_sql)
            await session.commit()
            return result.rowcount>0
        except Exception as e:
            await session.rollback()
            return False
        finally:
            await session.close()

    @staticmethod
    async def update_image_vector_by_id(
            vector_id:str,
            new_image_desc:str,
            new_desc_vector:list
    ):
        session:AsyncSession = AsyncSessionLocal()
        try:
            # 查询要更新的记录
            stmt = select(ImageVector).where(ImageVector.vector_id==vector_id)
            result = await  session.execute(stmt)
            image_vector_obj = result.scalar_one_or_none()

            if not image_vector_obj:
                print(f"未找到vector_id为{vector_id}的记录")
                return False

            #更新字段
            image_vector_obj.image_desc = new_image_desc
            image_vector_obj.desc_vector = new_desc_vector

            #提交更改
            await session.commit()
            print(f"成功更新vector_id为{vector_id}的记录")
            return True
        except Exception as e:
            await session.rollback()
            print(f"更新失败{str(e)}")
            return False
        finally:
            await session.close()



